import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm

class ConvSN(nn.Module):
    def __init__(self,in_ch,out_ch,kernel_size =4,padding=0,stride=2,pad_mode='zero',use_bias=True,sn=False):
        super(ConvSN,self).__init__()
        pad_layer = {
            'zero':nn.ZeroPad2d,
            'same':nn.ReplicationPad2d,
            'reflect':nn.ReflectionPad2d,
        }
        
        if pad_mode not in pad_layer:
            raise NotImplementedError
        
        self.pad_layer = pad_layer[pad_mode](padding)
        if sn:
            self.conv = spectral_norm(nn.Conv2d(in_ch,out_ch,kernel_size,stride,padding=0,bias=use_bias))
        else:
            self.conv = nn.Conv2d(in_ch,out_ch,kernel_size,stride,padding=0,bias=use_bias)

    def forward(self,input):
        out = self.pad_layer(input)
        out = self.conv(out)
        
        return out

class AnimeDiscriminator(nn.Module):
    def __init__(self, in_ch=64, out_ch=3, n_dis=3, sn=True):
        super(AnimeDiscriminator, self).__init__()
        channel = in_ch // 2

        self.conv1 = ConvSN(3, channel, kernel_size=3, stride=1, padding=1, sn=sn, use_bias=False)
        self.lrelu1 = nn.LeakyReLU(0.2, inplace=True)

        self.listLayer = []

        in_channel = channel
        for i in range(1, n_dis):
            self.listLayer.append(ConvSN(in_channel, channel*2, kernel_size=3, stride=2, padding=1, sn=sn, use_bias=False))
            self.listLayer.append(nn.LeakyReLU(0.2, inplace=True))

            self.listLayer.append(ConvSN(channel*2, channel*4, kernel_size=3, stride=1, padding=1, sn=sn, use_bias=False))
            self.listLayer.append(nn.InstanceNorm2d(channel*4, affine=True))
            self.listLayer.append(nn.LeakyReLU(0.2, inplace=True))

            in_channel = channel * 4
            channel = channel * 2
        self.listLayer = nn.ModuleList(self.listLayer)
        self.conv2 = ConvSN(in_channel, channel*2, kernel_size=3, stride=1, padding=1, sn=sn, use_bias=False)
        self.layernorm = nn.InstanceNorm2d(channel*2, affine=True)
        self.lrelu2 = nn.LeakyReLU(0.2, inplace=True)

        self.conv3 = ConvSN(channel*2, out_ch, kernel_size=3, stride=1, padding=1, sn=sn, use_bias=False)

    def forward(self, input):
        out = self.conv1(input)
        out = self.lrelu1(out)
        for layer in self.listLayer:
            out = layer(out)
        out = self.conv2(out)
        out = self.layernorm(out)
        out = self.lrelu2(out)
        out = self.conv3(out)
        return out



class ConvNormLReLU(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1, pad_mode="reflect", groups=1, bias=False,c_=None,h_=None):
        super(ConvNormLReLU, self).__init__()
        pad_layer = {
            "zero": nn.ZeroPad2d,
            "same": nn.ReplicationPad2d,
            "reflect": nn.ReflectionPad2d,
        }
        if pad_mode not in pad_layer:
            raise NotImplementedError

        self.pad = pad_layer[pad_mode](padding)
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, stride=stride, padding=0, groups=groups, bias=bias)
        
        
        
        self.norm = nn.InstanceNorm2d(out_ch, affine=True)
        # self.norm = nn.LayerNorm([c_,h_,h_])
        self.relu = nn.LeakyReLU(0.2, inplace=True)

    def forward(self, input):
        out = self.pad(input)
        out = self.conv(out)
        out = self.norm(out)
        out = self.relu(out)
        return out

class InvertedResBlock(nn.Module):
    def __init__(self, in_ch, out_ch, expansion_ratio=2):
        super(InvertedResBlock, self).__init__()

        self.use_res_connect = in_ch == out_ch
        bottleneck = int(round(in_ch * expansion_ratio))
        layers = []
        if expansion_ratio != 1:
            layers.append(ConvNormLReLU(in_ch, bottleneck, kernel_size=1, padding=0,c_=bottleneck,h_=64))

        # dw
        layers.append(ConvNormLReLU(bottleneck, bottleneck, groups=bottleneck, bias=True,c_=bottleneck,h_=64))
        # pw
        layers.append(nn.Conv2d(bottleneck, out_ch, kernel_size=1, padding=0, bias=False))
        layers.append(nn.InstanceNorm2d(out_ch, affine=True))
        # layers.append(nn.LayerNorm([out_ch,64,64]))
        self.layers = nn.Sequential(*layers)

    def forward(self, input):
        out = self.layers(input)
        if self.use_res_connect:
            out = input + out
        return out

    
class AnimeGenerator(nn.Module):
    def __init__(self, ):
        super(AnimeGenerator,self).__init__()
        # 3  256 256
        self.block_a = nn.Sequential(
            ConvNormLReLU(3, 32, kernel_size=7, padding=3,c_=32,h_=256),
            ConvNormLReLU(32, 64, stride=2, padding=(0, 1, 0, 1),c_=64,h_=128),
            ConvNormLReLU(64, 64,c_=64,h_=128)
        )

        self.block_b = nn.Sequential(
            ConvNormLReLU(64, 128, stride=2, padding=(0, 1, 0, 1),c_=128,h_=64),
            ConvNormLReLU(128, 128,c_=128,h_=64)
        )

        self.block_c = nn.Sequential(
            ConvNormLReLU(128, 128,c_=128,h_=64),
            InvertedResBlock(128, 256, 2),
            InvertedResBlock(256, 256, 2),
            InvertedResBlock(256, 256, 2),
            InvertedResBlock(256, 256, 2),
            ConvNormLReLU(256, 128,c_=128,h_=64),
        )

        self.block_d = nn.Sequential(
            ConvNormLReLU(128, 128,c_=128,h_=128),
            ConvNormLReLU(128, 128,c_=128,h_=128)
        )

        self.block_e = nn.Sequential(
            ConvNormLReLU(128, 64,c_=64,h_=256),
            ConvNormLReLU(64, 64,c_=64,h_=256),
            ConvNormLReLU(64, 32, kernel_size=7, padding=3,c_=32,h_=256)
        )

        self.out_layer = nn.Sequential(
            nn.Conv2d(32, 3, kernel_size=1, stride=1, padding=0, bias=False),
            nn.Tanh()
        )

    def forward(self, input):
        out = self.block_a(input)
        half_size = out.size()[-2:]
        out = self.block_b(out)
        out = self.block_c(out)

        out = F.interpolate(out, half_size, mode="bilinear", align_corners=True)
        out = self.block_d(out)
        out = F.interpolate(out, input.size()[-2:], mode="bilinear", align_corners=True)
        out = self.block_e(out)

        out = self.out_layer(out)
        return out